import cv2
import numpy as np
import os

# 创建保存图片的文件夹
save_path = "chessBoard1280"

# 棋盘格参数
chessboard_size = (8, 11)   # 12行9列棋盘格的内角点：11行8列
square_size = 10.0          # 每个格子的物理尺寸，单位毫米（mm）

# 已经保存的图片计数
image_count = 0
max_images = 16  # 需要保存的总张数

# 初始化摄像头
cap = cv2.VideoCapture(0)  # 0 表示默认摄像头

if not cap.isOpened():
    print("错误：无法打开摄像头")
    exit()

# 设置摄像头分辨率（关键修复！）
#cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
#cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

# 创建可调整大小的窗口（关键修复！）
cv2.namedWindow('Camera Calibration - 按"s"保存', cv2.WINDOW_NORMAL)
cv2.resizeWindow('Camera Calibration - 按"s"保存', 800, 600)  # 设置窗口大小

print("按键说明:")
print("  's' - 保存当前检测到棋盘格的帧")
print("  'q' - 退出程序")
print(f"需要采集 {max_images} 张图片")
print("请确保棋盘格在摄像头前，并调整角度和位置...")

while True:
    # 读取一帧
    ret, frame = cap.read()
    if not ret:
        print("错误：无法读取帧")
        break
    
    # 转换为灰度图（角点检测需要）
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    
    # 查找棋盘格角点
    ret_corners, corners = cv2.findChessboardCorners(gray, chessboard_size, None)
    
    # 如果找到角点，进行精细化并绘制
    if ret_corners:
        # 精细化角点坐标（达到亚像素精度）
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        corners_refined = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        
        # 在彩色图像上绘制角点和连线（彩色）
        frame_with_corners = cv2.drawChessboardCorners(frame.copy(), chessboard_size, corners_refined, ret_corners)
        
        # 在图像上显示提示信息
        cv2.putText(frame_with_corners, f"Detected! Press 's' to save ({image_count}/{max_images})", 
                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        # 显示带检测结果的帧
        cv2.imshow('Camera Calibration - 按"s"保存', frame_with_corners)
        
    else:
        # 没有检测到棋盘格，显示原始帧和提示
        cv2.putText(frame, f"Chessboard not found ({image_count}/{max_images})", 
                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        cv2.imshow('Camera Calibration - 按"s"保存', frame)
    
    # 等待按键输入
    key = cv2.waitKey(1) & 0xFF
    
    if key == ord('s') and ret_corners:  # 按下's'且检测到棋盘格时保存
        # 生成文件名
        filename = os.path.join(save_path, f"calibration_{image_count:02d}.jpg")
        # 保存原始帧（不是绘制了角点的帧），用于后续标定计算
        cv2.imwrite(filename, frame)
        image_count += 1
        print(f"已保存图片 {image_count}/{max_images}: {filename}")
        
        if image_count >= max_images:
            print(f"已完成 {max_images} 张图片采集！")
            break
            
    elif key == ord('q'):  # 按下'q'退出
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()

print(f"程序结束。共保存了 {image_count} 张图片在 '{save_path}' 目录中。")
print("接下来您可以使用这些图片进行相机标定计算。")

# 显示保存的图片数量统计
if image_count > 0:
    print("\n保存的图片列表:")
    for i in range(image_count):
        filename = f"calibration_{i:02d}.jpg"
        if os.path.exists(os.path.join(save_path, filename)):
            print(f"  {filename}")